# -*- coding: utf-8 -*-
"""
Created on Mon Mar 24 19:56:34 2025

@author: lenovo
"""

import pandas as pd
import numpy as np
from datetime import datetime,timedelta
from global_land_mask import globe
import re
from numba import njit
import matplotlib.pyplot as plt
import os
from shapely.geometry import Point, shape
import shapefile

# For basemap plot
from cartopy.mpl.gridliner import LongitudeFormatter, LatitudeFormatter
import matplotlib.ticker as mticker
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.geoaxes import GeoAxes
import cartopy.io.shapereader as shpreader
import geopandas as gpd


def split_str_id(names):
    strings = []
    numbers = []
    names = [names] if not isinstance(names, list) else names
    for name in names:
        parts = name.split('_')
        if len(parts) == 2:
            strings.append(parts[0])
            numbers.append(parts[1])  # 保留为字符串，不转 int
        else:
            raise ValueError("格式错误")
    return strings, numbers



def average_datetime(datetimes):
    if len(datetimes)==0:
        return None
    # 将datetime对象转换为时间戳并求和
    timestamps = [dt.timestamp() for dt in datetimes]
    avg_timestamp = sum(timestamps) / len(timestamps)
    # 将平均时间戳转换回datetime对象
    return datetime.fromtimestamp(avg_timestamp)


def see(data):
    # os.system('rm *png')
    if data.ndim==1:
        plt.plot(data)
        plt.savefig("./see.png")
        plt.show()
    elif data.ndim == 2:
        plt.contourf(data)
        plt.colorbar()
        plt.savefig("./see.png")
        plt.show()
    elif data.ndim == 3:
        num_slices = data.shape[0]  # 第一个维度的大小
        for i in range(num_slices):
            plt.figure()
            plt.contourf(data[i, :, :], cmap='viridis')  # 绘制第i个切片
            plt.colorbar()
            plt.savefig(f"./see{i}.png")
            plt.show()
    elif data.ndim == 4:
        num_slices = data.shape[1]  # 第二个维度的大小
        for i in range(num_slices):
            plt.figure()
            plt.contourf(data[0, i, :, :], cmap='viridis')  # 绘制第i个切片
            plt.colorbar()
            plt.savefig(f"./see{i}.png")
            plt.show()
    else:
        raise ValueError("输入数据的维度错误！")
    


class tydat:
    def __init__(self, path, RIstd, skiprows=2, lon_format='180'):
        """
        初始化台风数据
        
        Parameters:
        path: 数据文件路径
        RIstd: RI标准
        skiprows: 跳过的行数
        lon_format: 经度格式，'360' 表示0-360，'180' 表示-180到180
        """
        self.df = pd.read_csv(path, skiprows=skiprows, engine='python', sep=' & |\s', header=None)
        self.colors = ['#000000', '#00ffff', '#0000ff', '#FF8C00', '#FF0000', '#FF00FF']
        
        # 获取原始经纬度数据
        lon_raw = self.df[8].values
        lat_raw = self.df[9].values
        

        
        # 异常值筛选（先用原始数据筛选）
        self.control = self.cntrl_lonlat(lon_raw, lat_raw)
        
        # 获取筛选后的数据
        lon_filtered = lon_raw[self.control]
        lat_filtered = lat_raw[self.control]
        

        
        # 🔧 根据指定格式转换经度
        if lon_format == '360':
            # 转换为0-360格式
            self.lon = lon_filtered % 360

            
        elif lon_format == '180':
            # 转换为-180到180格式（原有逻辑）
            self.lon = np.where(lon_filtered > 180, lon_filtered - 360, lon_filtered)

            
        else:
            # 保持原始数据
            self.lon = lon_filtered

        
        # 其他属性
        self.lat = lat_filtered
        self.tp = self.df[14].values[self.control]
        self.umax = self.df[13].values[self.control]
        self.pmin = self.df[10].values[self.control]
        
        # 🔧 计算extent（考虑经度格式）
        dd = 8  # 调整边界缓冲
        
        if self.control.any():
            lon_min, lon_max = np.min(self.lon), np.max(self.lon)
            lat_min, lat_max = np.min(self.lat), np.max(self.lat)
            
            if lon_format == '360':
                # 检查是否跨越0度经线
                if lon_max - lon_min > 180:
                    print("⚠️ 检测到可能跨越0度经线的轨迹")
                    # 对于跨越0度的情况，可能需要特殊处理
                    # 这里先用简单的方法
                    
                self.extent = [lon_min - dd, lon_max + dd, lat_min - dd, lat_max + dd]
            else:
                self.extent = [lon_min - dd, lon_max + dd, lat_min - dd, lat_max + dd]
        else:
            self.extent = None
        
        
        self.RIstd = RIstd
        self.lon_format = lon_format
        
        # 处理时间数据
        try:
            self.time = np.array([datetime.strptime(str(t), '%Y%m%d%H') for t in self.df[0]])[self.control]
        except Exception as e:
            print(f"⚠️ 时间转换出现问题: {e}")
            self.time = None

    
    def cntrl_lonlat(self, lon, lat):
        """
        控制经纬度范围的异常值筛选
        """
        # 🔧 扩大经度范围检查，支持0-360格式
        lon_valid = (lon >= -180) & (lon <= 360)  # 支持更大范围
        lat_valid = (lat >= -90) & (lat <= 90)
        
        # 去除明显的异常值（如科学计数法错误）
        lon_reasonable = np.abs(lon) < 1e10
        lat_reasonable = np.abs(lat) < 1e10
        
        control = lon_valid & lat_valid & lon_reasonable & lat_reasonable
        
        if not control.any():
            None
        else:
            filtered_count = np.sum(control)
            total_count = len(control)

        
        return control
    
    def to_360_format(self):
        """转换为0-360格式"""
        if hasattr(self, 'lon_format') and self.lon_format == '360':
            return self
            
        
        # 转换经度
        self.lon = self.lon % 360
        
        # 更新extent
        if self.extent is not None:
            dd = 8
            lon_min, lon_max = np.min(self.lon), np.max(self.lon)
            self.extent[0] = lon_min - dd
            self.extent[1] = lon_max + dd
            

        
        self.lon_format = '360'

        
        return self
    
    def to_180_format(self):
        """转换为-180到180格式"""
        if hasattr(self, 'lon_format') and self.lon_format == '180':

            return self
            
        
        # 转换经度
        self.lon = ((self.lon + 180) % 360) - 180
        
        # 更新extent
        if self.extent is not None:
            dd = 8
            lon_min, lon_max = np.min(self.lon), np.max(self.lon)
            self.extent[0] = lon_min - dd
            self.extent[1] = lon_max + dd
        
        self.lon_format = '180'
        
        return self
    

    def num_rapidgrow(self):
        # count moment
        num=np.zeros(len(self.time))
        if self.RIstd == 15:
            for t in range(2,len(self.time)-2):
                gap = self.umax[t+2] - self.umax[t-2] 
                if gap>=15:
                    num[t]=1
        elif self.RIstd == 7:
            for t in range(1,len(self.time)-1):
                gap=self.umax[t+1] - self.umax[t-1]
                if gap>=7:
                    num[t]=1
        # print("single moment\n",num)
        #处理间断情况
        for t in range(len(num)):
            if t==0:
                pass
            if t>=1 and t<=len(num)-3 :
                if num[t]==0 and num[t-1]==1:
                    if num[t+1]==1 or num[t+2]==1:
                        num[t]=1
        # print("处理间断后\n",num)
        
        #处理陆地变形加强点;陆地取0
        num[globe.is_land(self.lat,self.lon)]=0
        # print("处理陆地后 \n",num)
        #处理孤立快增点
        for t in range(len(num)):
            if t==0 and num[t]==1:
                if num[t+2]+num[t+1]==2:
                    pass
                else:
                    num[t]==0
            if num[t-1]==0 and num[t]==1:
                if num[t+1]==0 :
                    num[t]=0        
        # print("处理孤立点后 \n",num)
        
        
        #moment串取中作为事件
        ##得到起始点
        start=[] ; end=[]
        for t in range(len(num)):
            if t==0 and num[t]==1:
                start.append(t)
            if t>0 and t<len(num)-1 : 
                if num[t]==1 and num[t-1]==0:
                    start.append(t)
                if num[t]==1 and num[t+1]==0:
                    end.append(t)
            if t==len(num)-1:
                if num[t]==1:
                    end.append(t)
        assert len(start)==len(end),f"len(start):{len(start),len(end):{len(end)}}"
        
        num = np.zeros(len(self.time))
        for i in range(len(start)):
            ave=int( (start[i]+end[i])/2 )
            num[ave]=1
        # print("取中后 \n",num)
        return num
    

def count_rapidgrow(RIstd,umax,t):
    '''
    umax:每个时次最大风速
    t:以datetime对象表示的所有时次
    '''
    from datetime import datetime,timedelta
    i=0
    record = np.zeros(len(umax))
    if RIstd == 15:
        '''RI=15的moment算法还没修改'''
        while t[i]+timedelta(days=1)<=t[-1]:
            trange = (t>=t[i])&(t<=t[i]+timedelta(days=1))
            urange = umax[trange]
            gap = max(urange)-urange[0]
            if gap>=RIstd:
                record[i]=1
            else:
                record[i]=0
            i+=1
    elif RIstd == 7:
        '''筛RI点时刻'''
        while t[i]+timedelta(hours=6)<=t[-1]:
            if t[i]-timedelta(hours=6)<t[0]:
                i+=1
                continue
            trange_forward  = (t<=t[i]+timedelta(hours=6)) & (t>=t[i])
            trange_backward = (t>=t[i]-timedelta(hours=6)) & (t<=t[i])
            maxumax_forward = max( umax[trange_forward] )
            maxumax_backward = min( umax[trange_backward] )
            gap = maxumax_forward - maxumax_backward
            if gap>=RIstd:
                record[i]=1
            else:
                record[i]=0
            i+=1
    num = record
    # print("single moment\n",num)
    #处理间断情况
    for t in range(len(num)):
        if t==0:
            pass
        if t>=1 and t<=len(num)-2-1 :
            if num[t]==0 and num[t-1]==1:
                if num[t+1]==1 or num[t+2]==1:
                    num[t]=1
    # print("处理间断后\n",num)
    
    #处理孤立快增点
    for t in range(len(num)):
        if t==0 and num[t]==1:
            if num[t+2]+num[t+1]==2:
                pass
            else:
                num[t]==0
        if num[t-1]==0 and num[t]==1:
            if num[t+1]==0 :
                num[t]=0        
    # print("处理孤立点后 \n",num)
    
    # obs一般无陆地RI点
    
    #moment串取中作为事件
    ##得到起始点
    start=[] ; end=[]
    for t in range(len(num)):
        if t==0 and num[t]==1:
            start.append(t)
        if t>0 and t<len(num)-1 : 
            if num[t]==1 and num[t-1]==0:
                start.append(t)
            if num[t]==0 and num[t-1]==1:
                end.append(t)
        if t==len(num)-1:
            if num[t]==1:
                end.append(t)
    assert len(start)==len(end),f"len(start):{len(start),len(end):{len(end)}}"
    
    num = np.zeros(len(umax))
    for i in range(len(start)):
        ave=int( (start[i]+end[i])/2 )
        num[ave]=1
    # print("取中后 \n",num)
    
    return num



class tydat_CMA:
    def __init__(self, path, skiprows=2, lon_format='180', timemode='6h'):
        """
        初始化CMA台风数据
        
        Args:
            path: 数据文件路径
            skiprows: 跳过的行数
            lon_format: 经度格式，'360' 表示0-360，'180' 表示-180到180
            timemode: 3h / 6h  3h for original observation; 6h for standard leadtime to align Ensemble

        """
        # 读取数据
        dfs = pd.read_csv(path, engine='python', sep='\s+', header=None, dtype=str)
        tindex = []
        time = []
        
        # 处理时间数据
        for i in range(len(dfs[0])):    
            t = dfs[0][i] + dfs[1][i] + dfs[2][i] + dfs[3][i]
            t = datetime.strptime(t, "%Y%m%d%H") - timedelta(hours=8)
            time.append(t)
            if np.isin(t.hour, [0, 6, 12, 18]).item():
                tindex.append(i)
        
        df = pd.read_csv(path, engine='python', sep='\s+', header=None)
        
        # 获取原始数据
        lon_raw = np.array(df[5], dtype=float)  # 确保是数值类型
        lat_raw = np.array(df[6], dtype=float)
        
        
        # 🔧 转换经度格式
        if lon_format == '360':
            # 转换为0-360格式
            lon_converted = lon_raw % 360
            lon2_converted = np.array([lon_converted[i] for i in tindex])
            
            
        elif lon_format == '180':
            # 保持-180到180格式
            lon_converted = ((lon_raw + 180) % 360) - 180
            lon2_converted = np.array([lon_converted[i] for i in tindex])
            
        else:
            # 不转换，使用原始数据
            lon_converted = lon_raw
            lon2_converted = np.array([lon_raw[i] for i in tindex])
        
        # 标准时刻数据（0, 6, 12, 18时）
        self.lon2 = lon2_converted
        self.lat2 = np.array([df[6][i] for i in tindex], dtype=float)
        self.umax2 = np.array([df[8][i] for i in tindex], dtype=float)
        self.pmin2 = np.array([df[7][i] for i in tindex], dtype=float)
        self.time2 = np.array([time[i] for i in tindex])
        
        # 全部时刻数据
        self.lon = lon_converted
        self.lat = lat_raw
        self.umax = np.array(df[8], dtype=float)
        self.pmin = np.array(df[7], dtype=float)
        self.time = np.array(time)
        
        # 🔧 计算extent（考虑经度格式）
        dd = 2  # 稍微调大边界缓冲
        
        if lon_format == '360':
            # 对于0-360格式，需要特别处理跨越0度的情况
            lon_min = np.min(self.lon)
            lon_max = np.max(self.lon)
            
            # 检查是否跨越0度经线
            if lon_max - lon_min > 180:
                self.extent = [lon_min - dd, lon_max + dd, 
                              np.min(self.lat) - dd, np.max(self.lat) + dd]
            else:
                self.extent = [lon_min - dd, lon_max + dd,
                              np.min(self.lat) - dd, np.max(self.lat) + dd]
        else:
            # 传统格式
            self.extent = [np.min(self.lon) - dd, np.max(self.lon) + dd,
                          np.min(self.lat) - dd, np.max(self.lat) + dd]
        
        # 存储格式信息
        self.lon_format = lon_format
        
        # 根据timemode 修改time和mode的值
        if timemode == '3h':
            pass 
        elif timemode=='6h':
            self.time = self.time2
            self.lon = self.lon2
            self.lat = self.lat2
            self.umax = self.umax2
            self.pmin = self.pmin2
        else:
            raise ValueError(f"timemode must be '3h' or '6h', but got {timemode}")




        
class tydat_JMA:
    def __init__(self,path,skiprows=1,extent=[100,140,10,30]):  # default - dusurui
        self.df=pd.read_csv(path,skiprows=skiprows,engine='python',header=None,sep="\s+") 
        self.colors=['#000000','#00ffff','#0000ff', '#FF8C00','#FF0000','#FF00FF']
        lat= self.df[3].values*0.1
        lon=self.df[4].values*0.1
        self.lon,self.lat=self.cntrl_lonlat(lon,lat)
        self.extent=extent
        # 这里需要转成datetime对象
        self.time = [datetime.strptime("20"+str(t),'%Y%m%d%H') for t in self.df[0] ]
        self.pmin = self.df[5]
        self.umax = self.df[6]*0.514  #kt to m/s
        
    def cntrl_lonlat(self,lon,lat):
        lon = lon[(lon>=0) & (lon <=360)] #去异常   有e25次方
        lon=np.where(lon>180,180,lon)  #防止左右边界相连
        return lon,lat
    
    
    
def extract_track_data(file_path, track_id):
    """
    从台风轨迹数据文件中提取指定 TRACK_ID 的数据，并将其分解为五个一维数组
    
    参数:
        file_path (str): 数据文件的路径
        track_id (str or int): 要提取的轨迹ID
    
    返回:
        tuple: 包含五个一维数组的元组 (time, lon, lat, pmin, vmax)
              如果找不到指定的轨迹ID，则返回五个空列表
    """
    import numpy as np
    # 将track_id转换为字符串以确保比较一致性
    track_id = str(track_id)
    
    # 初始化结果数组
    time = []
    lon = []
    lat = []
    pmin = []
    vmax = []
    
    # 读取文件
    with open(file_path, 'r') as file:
        lines = file.readlines()
    
    # 标记是否找到目标轨迹
    found_track = False
    # 标记是否正在处理目标轨迹的数据
    processing_target_track = False
    
    for line in lines:
        line = line.strip()
        
        # 跳过空行
        if not line:
            continue
        
        # 检查是否是轨迹ID行
        if line.startswith("TRACK_ID"):
            # 提取当前行的轨迹ID
            current_track_id = line.split()[1]
            
            # 如果找到目标轨迹ID
            if current_track_id == track_id:
                found_track = True
                processing_target_track = True
            else:
                # 如果之前正在处理目标轨迹，但现在遇到了新的轨迹ID，表示目标轨迹数据已处理完
                if processing_target_track:
                    break
                processing_target_track = False
        
        # 如果正在处理目标轨迹，且当前行不是轨迹ID、POINT_NUM、TRACK_NUM或标题行
        elif processing_target_track and not line.startswith(("TRACK_ID", "POINT_NUM", "TRACK_NUM", "0 0")):
            # 分割行数据
            parts = line.split()
            
            # 确保行包含足够的数据
            if len(parts) >= 6:
                time.append(parts[0])  # 时间
                lon.append(float(parts[1]))  # 经度
                lat.append(float(parts[2]))  # 纬度
                
                # 处理最小气压 (pmin)
                pmin_value = parts[3]
                if pmin_value == '1.0e+25':
                    pmin.append(float('nan'))  # 使用NaN表示缺失值
                else:
                    pmin.append(float(pmin_value))
                
                # 处理最大风速 (vmax)
                vmax_value = parts[5]
                if vmax_value == '1.0e+25':
                    vmax.append(float('nan'))  # 使用NaN表示缺失值
                else:
                    vmax.append(float(vmax_value))
    
    # 如果没有找到目标轨迹
    if not found_track:
        print(f"未找到 TRACK_ID {track_id}")
        
        return np.array([[]]*5)
    
    return np.array(time), np.array(lon), np.array(lat), np.array(pmin), np.array(vmax)


class tydat_NH():
    def __init__(self,file_path=r'D:\met_data\ty_obs\storms\storms_NH2023',track_id=5):
        '''filepath 指源文件地址；r'D:\met_data\ty_obs\storms\storms_NH2024'
        track_id = 5/8 
        '''
        time, lon, self.lat, self.pmin, self.umax = extract_track_data(file_path, track_id )
        self.time = np.array([datetime.strptime(date_string, '%Y%m%d%H') for date_string in time])
        self.lon = np.where(lon<=180,lon,lon-360)
        dd = 1
        self.extent=[np.min(self.lon)-dd,np.max(self.lon)+dd,np.min(self.lat)-dd,np.max(self.lat)+dd]
        
        
        
def load_land_polygons(shp_path: str, dbf_path: str):
    """
    读取 Natural Earth 陆地 shapefile，返回 shapely Polygon 列表。
    """
    reader = shapefile.Reader(shp=shp_path, dbf=dbf_path)
    return [shape(rec) for rec in reader.shapes()]



def detect_landfall(lats, lons, land_polygons):
    """
    判断一段台风轨迹是否登陆。
    参数：
        lats: 可迭代的纬度序列（浮点数，单位°）
        lons: 可迭代的经度序列（浮点数，单位°），若 lon>180，会自动转换为 –180…180 范围
        land_polygons: load_land_polygons() 返回的多边形列表
    返回：
        flags: 与输入点对应的布尔列表，True 表示该点“在陆地上”
        first_landfall_idx: 第一次登陆的索引（flags 中第一个 True 的位置），
                            若全程未登陆则返回 None
    """
    flags = []
    for lat, lon in zip(lats, lons):
        # 归一化经度
        if lon > 180:
            lon -= 360.0
        pt = Point(lon, lat)
        # 只要落入任意一个陆地多边形就算登陆
        on_land = any(pt.within(poly) for poly in land_polygons)
        flags.append(on_land)

    # 找到第一次登陆的点
    try:
        first_landfall_idx = flags.index(True)
    except ValueError:
        first_landfall_idx = None

    return flags, first_landfall_idx

    

def basemap(ax, title=None, extent=None, add_province=False, lon_interval=20, lat_interval=10):
    """
    只显示中国省界的底图
    
    Parameters:
    -----------
    ax : cartopy GeoAxes
        地图坐标轴
    title : str, optional
        图标题
    extent : list, optional
        地图范围 [lon_min, lon_max, lat_min, lat_max]
        默认: None (不设置范围)
    lon_interval : float, optional
        经度线间隔（默认5度）
    lat_interval : float, optional
        纬度线间隔（默认5度）
    """
    
    # ⭐ 检查是否是 GeoAxes（有 projection）
    if not isinstance(ax, GeoAxes):
        raise TypeError(
            "❌ Error: ax 必须是 cartopy GeoAxes！\n"
            "请在创建时添加 projection，例如：\n"
            "  ax = fig.add_subplot(111, projection=ccrs.PlateCarree())\n"
            "或者使用 create_map_axes() 函数自动创建"
        )
    
    if add_province:
        provinces_shp = shpreader.natural_earth(
                resolution='10m',
                category='cultural',
                name='admin_1_states_provinces'  # 面状，不是 _lines
            )
        gdf = gpd.read_file(provinces_shp)
        # 筛选中国省份
        china_provinces = gdf[
            (gdf['admin'] == 'China') | 
            (gdf['adm0_a3'] == 'CHN')
        ]
        # 绘制省界
        ax.add_geometries(
            china_provinces.geometry,
            ccrs.PlateCarree(),
            facecolor='none',
            edgecolor='black',
            linewidth=0.5
        )

    # 添加地图要素
    ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
    # ax.add_feature(cfeature.BORDERS, linewidth=0.5)
    
    # 设置地图范围
    if extent is not None:
        ax.set_extent(extent, crs=ccrs.PlateCarree())
    
    # ⭐ 网格线（关键修改）
    gl = ax.gridlines(
        draw_labels=True, 
        alpha=0.3, 
        linestyle='--', 
        linewidth=0.5,
        color='gray'
    )
    
    # ⭐⭐⭐ 设置经纬度刻度间隔
    gl.xlocator = mticker.MultipleLocator(lon_interval)  # 经度间隔
    gl.ylocator = mticker.MultipleLocator(lat_interval)  # 纬度间隔
    
    # 标签格式
    gl.top_labels = False
    gl.right_labels = False
    gl.xlabel_style = {'size': 8}
    gl.ylabel_style = {'size': 8}
    gl.xformatter = LongitudeFormatter(number_format='.1f')
    gl.yformatter = LatitudeFormatter(number_format='.1f')
    
    if title:
        ax.set_title(title, fontsize=12, fontweight='bold')


import pickle 

def save_pickle(data, save_path):
    """将数据保存为pickle文件
    参数:
        data: 要保存的数据
        save_path: 保存路径，需以.pkl结尾
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    # 保存为pickle文件
    with open(save_path, 'wb') as f:
        pickle.dump(data, f)


def load_pickle(load_path):
    """从pickle文件中加载数据
    参数:
        load_path: 加载路径，需以.pkl结尾
    返回:
        加载的数据
    """
    # 确保文件存在
    if not os.path.exists(load_path):
        raise FileNotFoundError(f"File not found: {load_path}")
    # 加载pickle文件
    with open(load_path, 'rb') as f:
        data = pickle.load(f)
    return data